#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
路径显示节点 - 在ComfyUI界面中显示路径信息和文件列表
"""

import os
import json
import folder_paths
from typing import List, Dict, Any, Optional
import comfy.utils

class PathInfoDisplayNode:
    """路径信息显示节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path": ("STRING", {"default": "", "multiline": False}),
                "show_details": ("BOOLEAN", {"default": True}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("path_info", "file_count", "folder_count")
    FUNCTION = "display_path_info"
    CATEGORY = "AI训练/path_manager"

    def display_path_info(self, path, show_details):
        """显示路径信息"""
        try:
            if not path or not os.path.exists(path):
                return ("路径不存在", "0", "0")
            
            path_info = {
                "path": path,
                "exists": True,
                "is_file": os.path.isfile(path),
                "is_dir": os.path.isdir(path),
                "is_absolute": os.path.isabs(path),
                "size": 0,
                "permissions": {}
            }
            
            if path_info["is_file"]:
                path_info["size"] = os.path.getsize(path)
                file_count = "1"
                folder_count = "0"
            elif path_info["is_dir"]:
                files = [f for f in os.listdir(path) if os.path.isfile(os.path.join(path, f))]
                folders = [f for f in os.listdir(path) if os.path.isdir(os.path.join(path, f))]
                file_count = str(len(files))
                folder_count = str(len(folders))
                
                if show_details:
                    path_info["files"] = files[:10]  # 只显示前10个文件
                    path_info["folders"] = folders[:10]  # 只显示前10个文件夹
                    path_info["total_files"] = len(files)
                    path_info["total_folders"] = len(folders)
            else:
                file_count = "0"
                folder_count = "0"
            
            # 获取权限信息
            path_info["permissions"] = {
                "readable": os.access(path, os.R_OK),
                "writable": os.access(path, os.W_OK),
                "executable": os.access(path, os.X_OK)
            }
            
            return (json.dumps(path_info, indent=2), file_count, folder_count)
            
        except Exception as e:
            return (f"错误: {str(e)}", "0", "0")

class FileListDisplayNode:
    """文件列表显示节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "folder_path": ("STRING", {"default": "", "multiline": False}),
                "file_filter": ("STRING", {"default": "*", "multiline": False}),
                "max_files": ("INT", {"default": 50, "min": 1, "max": 1000}),
                "sort_by": (["name", "size", "modified"],),
                "sort_order": (["asc", "desc"],),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "STRING")
    RETURN_NAMES = ("file_list", "file_info", "summary")
    FUNCTION = "display_file_list"
    CATEGORY = "AI训练/path_manager"

    def display_file_list(self, folder_path, file_filter, max_files, sort_by, sort_order):
        """显示文件列表"""
        try:
            if not folder_path or not os.path.exists(folder_path) or not os.path.isdir(folder_path):
                return ("[]", "{}", "文件夹不存在")
            
            files = []
            for item in os.listdir(folder_path):
                item_path = os.path.join(folder_path, item)
                if os.path.isfile(item_path) and self._matches_filter(item, file_filter):
                    file_info = {
                        "name": item,
                        "path": item_path,
                        "size": os.path.getsize(item_path),
                        "modified": os.path.getmtime(item_path),
                        "extension": os.path.splitext(item)[1].lower()
                    }
                    files.append(file_info)
            
            # 排序
            reverse = sort_order == "desc"
            if sort_by == "name":
                files.sort(key=lambda x: x["name"].lower(), reverse=reverse)
            elif sort_by == "size":
                files.sort(key=lambda x: x["size"], reverse=reverse)
            elif sort_by == "modified":
                files.sort(key=lambda x: x["modified"], reverse=reverse)
            
            # 限制文件数量
            files = files[:max_files]
            
            # 生成摘要
            total_size = sum(f["size"] for f in files)
            extensions = {}
            for f in files:
                ext = f["extension"]
                extensions[ext] = extensions.get(ext, 0) + 1
            
            summary = {
                "total_files": len(files),
                "total_size": total_size,
                "size_mb": round(total_size / (1024 * 1024), 2),
                "extensions": extensions
            }
            
            return (
                json.dumps([f["name"] for f in files]),
                json.dumps(files),
                json.dumps(summary)
            )
            
        except Exception as e:
            return ("[]", "{}", f"错误: {str(e)}")
    
    def _matches_filter(self, filename, file_filter):
        """检查文件是否匹配过滤器"""
        if file_filter == "*":
            return True
        return filename.lower().endswith(file_filter.lower())

class FolderTreeDisplayNode:
    """文件夹树形显示节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "root_path": ("STRING", {"default": "", "multiline": False}),
                "max_depth": ("INT", {"default": 3, "min": 1, "max": 10}),
                "include_files": ("BOOLEAN", {"default": True}),
                "file_filter": ("STRING", {"default": "*", "multiline": False}),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("folder_tree", "tree_info")
    FUNCTION = "display_folder_tree"
    CATEGORY = "AI训练/path_manager"

    def display_folder_tree(self, root_path, max_depth, include_files, file_filter):
        """显示文件夹树形结构"""
        try:
            if not root_path or not os.path.exists(root_path) or not os.path.isdir(root_path):
                return ("{}", "路径不存在")
            
            tree = self._build_tree(root_path, max_depth, include_files, file_filter, 0)
            
            tree_info = {
                "root": root_path,
                "max_depth": max_depth,
                "total_folders": self._count_folders(tree),
                "total_files": self._count_files(tree) if include_files else 0
            }
            
            return (json.dumps(tree, indent=2), json.dumps(tree_info))
            
        except Exception as e:
            return ("{}", f"错误: {str(e)}")
    
    def _build_tree(self, path, max_depth, include_files, file_filter, current_depth):
        """构建树形结构"""
        if current_depth > max_depth:
            return {}
        
        tree = {}
        
        try:
            for item in os.listdir(path):
                item_path = os.path.join(path, item)
                
                if os.path.isdir(item_path):
                    subtree = self._build_tree(item_path, max_depth, include_files, file_filter, current_depth + 1)
                    tree[item] = {"type": "folder", "content": subtree}
                elif include_files and os.path.isfile(item_path) and self._matches_filter(item, file_filter):
                    tree[item] = {"type": "file", "size": os.path.getsize(item_path)}
        except PermissionError:
            tree["[权限不足]"] = {"type": "error"}
        except Exception as e:
            tree[f"[错误: {str(e)}]"] = {"type": "error"}
        
        return tree
    
    def _matches_filter(self, filename, file_filter):
        """检查文件是否匹配过滤器"""
        if file_filter == "*":
            return True
        return filename.lower().endswith(file_filter.lower())
    
    def _count_folders(self, tree):
        """统计文件夹数量"""
        count = 0
        for item, info in tree.items():
            if info.get("type") == "folder":
                count += 1
                count += self._count_folders(info.get("content", {}))
        return count
    
    def _count_files(self, tree):
        """统计文件数量"""
        count = 0
        for item, info in tree.items():
            if info.get("type") == "file":
                count += 1
            elif info.get("type") == "folder":
                count += self._count_files(info.get("content", {}))
        return count

class PathComparisonNode:
    """路径比较节点"""
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "path1": ("STRING", {"default": "", "multiline": False}),
                "path2": ("STRING", {"default": "", "multiline": False}),
                "compare_type": (["existence", "content", "structure"],),
            }
        }
    
    RETURN_TYPES = ("STRING", "STRING", "BOOLEAN")
    RETURN_NAMES = ("comparison_result", "detailed_info", "are_equal")
    FUNCTION = "compare_paths"
    CATEGORY = "AI训练/path_manager"

    def compare_paths(self, path1, path2, compare_type):
        """比较两个路径"""
        try:
            result = {
                "path1": path1,
                "path2": path2,
                "compare_type": compare_type,
                "path1_exists": os.path.exists(path1),
                "path2_exists": os.path.exists(path2),
                "comparison": {}
            }
            
            are_equal = False
            
            if compare_type == "existence":
                are_equal = result["path1_exists"] == result["path2_exists"]
                result["comparison"] = {
                    "both_exist": result["path1_exists"] and result["path2_exists"],
                    "both_not_exist": not result["path1_exists"] and not result["path2_exists"],
                    "only_path1_exists": result["path1_exists"] and not result["path2_exists"],
                    "only_path2_exists": not result["path1_exists"] and result["path2_exists"]
                }
            
            elif compare_type == "content" and result["path1_exists"] and result["path2_exists"]:
                if os.path.isfile(path1) and os.path.isfile(path2):
                    # 比较文件内容
                    size1 = os.path.getsize(path1)
                    size2 = os.path.getsize(path2)
                    are_equal = size1 == size2
                    result["comparison"] = {
                        "size1": size1,
                        "size2": size2,
                        "size_equal": size1 == size2
                    }
                elif os.path.isdir(path1) and os.path.isdir(path2):
                    # 比较目录内容
                    files1 = set(os.listdir(path1))
                    files2 = set(os.listdir(path2))
                    common_files = files1 & files2
                    only_in_path1 = files1 - files2
                    only_in_path2 = files2 - files1
                    
                    are_equal = len(only_in_path1) == 0 and len(only_in_path2) == 0
                    result["comparison"] = {
                        "common_files": list(common_files),
                        "only_in_path1": list(only_in_path1),
                        "only_in_path2": list(only_in_path2),
                        "total_files_path1": len(files1),
                        "total_files_path2": len(files2)
                    }
            
            elif compare_type == "structure" and result["path1_exists"] and result["path2_exists"]:
                if os.path.isdir(path1) and os.path.isdir(path2):
                    # 比较目录结构
                    structure1 = self._get_structure(path1)
                    structure2 = self._get_structure(path2)
                    
                    are_equal = structure1 == structure2
                    result["comparison"] = {
                        "structure1": structure1,
                        "structure2": structure2,
                        "structure_equal": structure1 == structure2
                    }
            
            return (json.dumps(result), json.dumps(result["comparison"]), are_equal)
            
        except Exception as e:
            return (f"比较错误: {str(e)}", "{}", False)
    
    def _get_structure(self, path, max_depth=3, current_depth=0):
        """获取目录结构"""
        if current_depth > max_depth:
            return {}
        
        structure = {}
        try:
            for item in os.listdir(path):
                item_path = os.path.join(path, item)
                if os.path.isdir(item_path):
                    structure[item] = self._get_structure(item_path, max_depth, current_depth + 1)
                else:
                    structure[item] = "file"
        except:
            structure["[无法访问]"] = "error"
        
        return structure

# 节点映射
NODE_CLASS_MAPPINGS = {
    "PathInfoDisplayNode": PathInfoDisplayNode,
    "FileListDisplayNode": FileListDisplayNode,
    "FolderTreeDisplayNode": FolderTreeDisplayNode,
    "PathComparisonNode": PathComparisonNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PathInfoDisplayNode": "路径信息显示",
    "FileListDisplayNode": "文件列表显示",
    "FolderTreeDisplayNode": "文件夹树形显示",
    "PathComparisonNode": "路径比较器",
} 